import torch
import torch.nn as nn
from torch.nn import init

class ChannelAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.GELU(),
            nn.Conv2d(channel // reduction, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result = self.maxpool(x)
        avg_result = self.avgpool(x)
        max_out = self.se(max_result)
        avg_out = self.se(avg_result)
        out = self.sigmoid(max_out + avg_out)
        return out

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], dim=1)
        out = self.conv(result)
        out = self.sigmoid(out)
        return out

class CBAMBlock(nn.Module):
    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        self.ChannelAttention = ChannelAttention(channel=channel, reduction=reduction)
        self.SpatialAttention = SpatialAttention(kernel_size=kernel_size)
        self.norm = nn.BatchNorm2d(channel)

    def forward(self, x):
        # B, C, H, W = x.size()
        residual = x
        out = x * self.ChannelAttention(x)
        out = out * self.SpatialAttention(out)
        out = self.norm(out + residual)
        return out

if __name__ == '__main__':
    x = torch.rand(1, 192, 56, 56)
    model = CBAMBlock(channel=192, reduction=16, kernel_size=7)
    print(model(x).shape) #torch.Size([1, 192, 56, 56])